import logging
import numpy as np
import pinecone
from tensorflow.keras.datasets import mnist
from tqdm import tqdm
from collections import Counter
import matplotlib.pyplot as plt
from pinecone import Pinecone, ServerlessSpec

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

# 初始化 Pinecone
pinecone = Pinecone(api_key="67ad532a-d0c4-4aca-9b00-a7ea204cc919")
index_name = "mnist-index"

# 获取现有索引列表
existing_indexes = pinecone.list_indexes()

# 检查索引是否存在，如果存在就删除
if any(index['name'] == index_name for index in existing_indexes):
    logging.info(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    logging.info(f"索引 '{index_name}' 已成功删除。")

# 创建新索引
logging.info(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=784,  # MNIST 每个图像展平后是一个 784 维向量
    metric="euclidean",  # 使用欧氏距离
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
logging.info(f"索引 '{index_name}' 创建成功。")

# 连接到索引
index = pinecone.Index(index_name)
logging.info(f"已成功连接到索引 '{index_name}'。")

# 加载MNIST数据集
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# 添加一个小的epsilon值以避免零值
epsilon = 1e-8

# 扁平化、归一化并添加epsilon
x_train = x_train.reshape(-1, 28*28).astype('float32') / 255.0 + epsilon
x_test = x_test.reshape(-1, 28*28).astype('float32') / 255.0 + epsilon

num_samples = 1000
x_train = x_train[:num_samples]
y_train = y_train[:num_samples]

# 准备上传到Pinecone的数据
vectors = [
    (str(i), x_train[i].tolist(), {"label": int(y_train[i])})
    for i in range(len(x_train))
]

# 上传数据到Pinecone
batch_size = 100
max_retries = 3
logging.info("正在上传数据到Pinecone...")
for i in tqdm(range(0, len(vectors), batch_size), desc="上传数据", unit="it"):
    batch = vectors[i:i + batch_size]
    index.upsert(batch)
logging.info(f"成功上传了 {len(vectors)} 条记录。")

# 测试准确率函数
def test_accuracy(k):
    correct_predictions = 0
    for i in tqdm(range(len(x_test)), desc="测试准确率", unit="项"):
        query_result = index.query(
            vector=x_test[i].tolist(),  # 从 queries=[x_test[i].tolist()] 改为 vector=x_test[i].tolist()
            top_k=k,
            include_metadata=True
        )
        nearest_labels = [match['metadata']['label'] for match in query_result['matches']]
        if y_test[i] in nearest_labels:
            correct_predictions += 1
    accuracy = correct_predictions / len(x_test)
    return accuracy


# 使用k=11测试准确率
k_value = 11
accuracy = test_accuracy(k_value)
logging.info(f"使用Pinecone的准确率 (k={k_value}): {accuracy:.4f}")

# 清理：删除索引（可选）
pinecone.delete_index(index_name)